# @matching/first_stage.py
import json
import logging
import threading
import numpy as np
from bson import ObjectId
from pymongo import MongoClient
from utils.mongo_utils import get_collection
from config.config import Config, MatchType, MatchMethod
from matching.filter_by_name import get_npm_candidates_fuzzy, filter_candidates_by_name_similarity
from entity.function_candidate import FunctionCandidate
# 全局缓存（避免重复加载）
_faiss_index = None
_id_map = None

# 全局锁
_faiss_load_lock = threading.Lock()


def load_faiss_index():
    import faiss
    """懒加载 FAISS 索引和 ID 映射（线程安全）"""
    global _faiss_index, _id_map
    if _faiss_index is not None and _id_map is not None:
        return _faiss_index, _id_map

    with _faiss_load_lock:  # 加锁
        # 双重检查（避免锁内重复加载）
        if _faiss_index is not None and _id_map is not None:
            return _faiss_index, _id_map

        logging.debug("正在加载 FAISS 索引...")
        _faiss_index = faiss.read_index(Config.FAISS_INDEX_PATH)
        with open(Config.ID_MAP_PATH, 'rb') as f:
            _id_map = json.load(f)
        logging.debug(f"FAISS 索引加载完成，共 {len(_id_map)} 个向量")

    return _faiss_index, _id_map


def vector_similarity_stage_by_faiss(
        arkts_doc,
        arkts_id,
        base_query=None,
):
    import faiss
    """
    第一阶段：使用 FAISS 快速搜索高/中相似度候选函数
    """
    """
        第一阶段：使用 FAISS 快速搜索高/中相似度候选函数
        """
    # 1. 加载 FAISS 索引
    try:
        index, id_map = load_faiss_index()
    except Exception as e:
        logging.error(f"[向量][{arkts_id}] 加载 FAISS 索引失败: {e}")
        return [], [], False

    # 2. 获取查询向量
    arkts_vector = arkts_doc.get('vector')
    if not arkts_vector:
        logging.error(f"[向量][{arkts_id}] ⚠️ ArkTS函数没有向量表示，跳过筛选")
        return [], [], False

    if len(arkts_vector) != index.d:
        logging.error(f"[向量][{arkts_id}] 向量维度不匹配: 期望 {index.d}, 实际 {len(arkts_vector)}")
        return [], [], False

    # 3. 转为 numpy 并归一化（关键！）
    query_vec = np.array(arkts_vector, dtype='float32').reshape(1, -1)
    faiss.normalize_L2(query_vec)

    # 4. 设置 nprobe（可调）
    index.nprobe = 20  # 根据你的 nlist 设置（如 nlist=4000）

    # 5. 执行搜索（Top-1000 候选足够）
    k = 300
    distances, indices = index.search(query_vec, k)

    # distances[0]: 相似度数组（内积 = 余弦相似度）
    # indices[0]: FAISS 内部编号

    high_similarity_candidates = []
    medium_similarity_candidates = []

    # 6. 转换为 MongoDB _id 并过滤
    for i, faiss_idx in enumerate(indices[0]):
        if faiss_idx == -1:
            continue  # 没有足够候选
        # 关键修复：将NumPy int64转换为Python int，然后转为字符串
        faiss_idx_int = int(faiss_idx)  # 先转为Python int
        faiss_idx_str = str(faiss_idx_int)  # 再转为字符串

        # 检查键是否存在
        if faiss_idx_str not in id_map:
            logging.warning(f"FAISS索引 {faiss_idx_int} 在ID映射中不存在")
            continue
        similarity = float(distances[0][i])  # 同样转换distance
        try:
            npm_id = id_map[faiss_idx_str]  # 映射到 MongoDB _id
        except IndexError:
            continue

        # 可选：如果 base_query 有 _id $in 限制，检查是否在候选集中
        if base_query and '_id' in base_query and '$in' in base_query['_id']:
            object_id = ObjectId(npm_id)
            if object_id not in base_query['_id']['$in']:
                continue  # 不在名称匹配结果中，跳过


        candidate = FunctionCandidate(
            candidate_id_name="npm_id",
            candidate_id=npm_id,
            vector_similarity=similarity
        )
        if similarity >= 0.90:
            high_similarity_candidates.append(candidate)
        elif similarity >= 0.5:
            medium_similarity_candidates.append(candidate)

    # 7. 排序
    high_similarity_candidates.sort(key=lambda x: x.vector_similarity, reverse=True)
    medium_similarity_candidates.sort(key=lambda x: x.vector_similarity, reverse=True)
    logging.debug(
        f"[向量][{arkts_id}] FAISS 搜索完成: 高相似 {len(high_similarity_candidates)}, "
        f"中相似 {len(medium_similarity_candidates)}"
    )

    return high_similarity_candidates, medium_similarity_candidates, True


#单纯对待检测的向量进行余弦相似度计算，这个计算TPL Version时候用到的
def vector_similarity_stage(npm_doc, npm_id, base_query=None, batch_size=5000):
    # 1. 初始化检查
    npm_vector = npm_doc.get('vector')
    if not npm_vector:
        logging.error(f"[向量][{npm_id}] ⚠️NPM函数没有向量表示，跳过筛选")
        return [], [], False

    # 2. 预处理目标向量
    npm_vec = np.array(npm_vector, dtype=np.float32)
    npm_norm = np.linalg.norm(npm_vec)
    if npm_norm == 0:
        return [], [], False
    npm_vec_normalized = npm_vec / npm_norm

    # 3. 准备数据库查询
    arkts_functions = get_collection(Config.DB_NAME, Config.ArkTS_COLLECTION)
    projection = {"_id": 1, "vector": 1}
    query = base_query.copy() if base_query else {}

    # 4. 初始化结果和分页
    high_sim_candidates = []
    medium_sim_candidates = []
    skip = 0

    while True:
        # 4.1 执行分页查询（保持原skip逻辑）
        batch = list(arkts_functions.find(query, projection).skip(skip).limit(batch_size))
        if not batch:
            break

        # 4.2 更新分页标记
        skip += len(batch)

        # 4.3 处理当前批次
        valid_ids = []
        valid_vectors = []
        for doc in batch:
            if doc.get("vector"):
                valid_ids.append(doc["_id"])
                valid_vectors.append(doc["vector"])

        # 如果没有有效向量，跳过计算
        if not valid_vectors:
            continue

        # 5. 向量计算
        vectors = np.array(valid_vectors, dtype=np.float32)
        # 归一化向量矩阵
        norms = np.linalg.norm(vectors, axis=1, keepdims=True)
        # 避免除零（将零向量归一化为零向量）
        norms[norms == 0] = 1
        normalized_vectors = vectors / norms

        # 计算相似度
        similarities = np.dot(normalized_vectors, npm_vec_normalized)

        # 6. 分类处理结果
        for i, sim in enumerate(similarities):
            candidate = FunctionCandidate(
                candidate_id_name="arkts_id",
                candidate_id=str(valid_ids[i]),
                vector_similarity=float(sim)
            )
            if sim > 0.90:
                candidate.match_method = MatchMethod.VECTOR_ONLY.value
                high_sim_candidates.append(candidate)
            elif sim > 0.4:
                medium_sim_candidates.append(candidate)
    # 7. 排序
    high_sim_candidates.sort(key=lambda x: x.vector_similarity, reverse=True)
    medium_sim_candidates.sort(key=lambda x: x.vector_similarity, reverse=True)
    logging.debug(
        f"[向量][{npm_id}] 全集 搜索完成: 高相似 {len(high_sim_candidates)}, "
        f"中相似 {len(medium_sim_candidates)}"
    )
    return high_sim_candidates, medium_sim_candidates, True


def first_stage_filter(target_doc, target_id, match_type,
                       ark_ts_library = None,
                       require_name_match=False):
    """
    第一阶段筛选函数
    :param ark_ts_library: 用于TPL Version匹配时候用到
    :param match_type: 匹配的过程类型，TPL匹配 和 TPLVersion
    :param target_id:ArkTS 函数id(TPL 匹配时候） ，NPM函数id（TPL Version匹配时候）
    :param target_doc: ArkTS 函数文档(TPL 匹配时候） ，NPM函数文档（TPL Version匹配时候）
    :param require_name_match: 是否使用名称匹配
    """
    # --- 新增：如果需要名称匹配，则执行模糊匹配逻辑 ---
    if require_name_match:
        arkts_name = target_doc.get('name', '').strip()
        if arkts_name and arkts_name.lower() != 'anonymous':
            logging.debug(f"开始为 待检测 函数 '{arkts_name}' 进行模糊名称匹配...")
            # 1. 获取初步候选 (使用 name_filter.py 的查询逻辑)
            npm_candidates = get_npm_candidates_fuzzy(arkts_name)
            logging.debug(f"  初步获取到 {len(npm_candidates)} 个候选函数。")

            # 2. 使用 fuzzywuzzy 进行二次过滤
            filtered_candidate_ids = filter_candidates_by_name_similarity(arkts_name, npm_candidates)
            logging.debug(f"  模糊匹配后筛选出 {len(filtered_candidate_ids)} 个候选函数。")

            # 如果筛选后没有候选，则直接返回空
            if not filtered_candidate_ids:
                logging.debug("  模糊匹配后无候选函数，跳过向量/哈希筛选。")
                return [], [], False  # 注意：这里 has_vector_or_hash 为 False 是合理的，因为没进行计算

            # --- 修改：构建一个只包含这些 ID 的查询，传递给后续阶段 ---
            base_query = {
                'ast_feature': {'$exists': True},
                '_id': {'$in': [ObjectId(cid) for cid in filtered_candidate_ids]}
            }
        else:
            # 如果是匿名函数或无名函数，则不应用名称过滤，base_query 回退到默认
            logging.debug("ArkTS 函数为匿名或无名，不应用名称匹配。")
            base_query = {'ast_feature': {'$exists': True}}
    else:
        # 如果不要求名称匹配，则使用默认查询
        base_query = {'ast_feature': {'$exists': True}}
    if match_type == MatchType.TPL_MATCH:
        # 将 base_query 传递给向量阶段
        return vector_similarity_stage_by_faiss(target_doc, target_id,
                                                base_query=base_query)
    elif match_type == MatchType.TPL_VERSION_MATCH:
        if ark_ts_library:
            base_query['library'] = ark_ts_library # 这里使用library查询的更加快速一些
        return vector_similarity_stage(target_doc, target_id,base_query)
    else:
        logging.error("没有包含正确的预处理命令")
        return [], [], False  # 返回空结果和 False 标志
